from keras import models
import numpy as np
from imageDeal import ImageDeal

# 使用模型进行预测
class HandWrittenDigitRecognizer(object):
    # 传入模型路径 与 图片集合
    def __init__(self, modelPath) -> None:
        self.modelPath = modelPath
        self.model = None

        self._loadModel()

    # 导入模型
    def _loadModel(self):
        self.model = models.load_model(self.modelPath)

    # 获取标签集合
    def getLabelList(self, targetImageList):
        # 归一化
        targetImageList = targetImageList.astype("float32") / 255
        # 前置标签
        preLabelList = self.model.predict(targetImageList)

        labelList = []
        for preLabel in preLabelList:
            # 结合 softmax 函数的输出使用
            labelList.append(np.argmax(preLabel))

        return labelList

if __name__ == "__main__":
    modelPath = "./model/model.h5"

    recognizer = HandWrittenDigitRecognizer(modelPath)

    ID = ImageDeal("./image/test.png")
    imageList = ID.transMNIST(ID.getImageSagList())

    labelList = recognizer.getLabelList(imageList)

    print("the result of label: {}".format(labelList))
